#!/usr/bin/env python3
import sys
import argparse
import re
import cx_Oracle
import requests
import yaml
from datetime import datetime
from tenacity import retry, wait_exponential, stop_after_attempt
from typing import Tuple, Iterator, Optional

# 在parse_args中新增配置文件参数
def parse_args():
    parser = argparse.ArgumentParser(description='数据库实例状态检查工具')
    parser.add_argument('-g', '--group', help='精确匹配组名称')
    parser.add_argument('-c', '--config', required=True, help='数据库配置文件路径')
    parser.add_argument('--alert-config', required=True, help='告警配置文件路径')
    parser.add_argument('--wx-retry', type=int, default=3, help='企业微信通知重试次数')
    return parser.parse_args()

# 修改主逻辑中的配置加载方式


# @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, max=10))
def check_instance(db_name: str, conn_str: str) -> str:
    try:
        # 直接创建连接
        with cx_Oracle.connect(user=user, password=pwd, dsn=dsn) as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT status FROM v$instance")
                return cur.fetchone()[0]
    except Exception as e:
        print(f"连接尝试参数: {conn_str}")
        raise RuntimeError(f"数据库检查失败: {str(e)}")

class AlertConfig:
    def __init__(self, webhook: str, agentid: str):
        self.webhook = webhook
        self.agentid = agentid

class DBConnectionConfig:
    def __init__(self, db_name: str, user: str, password: str, dsn: str):
        self.db_name = db_name
        self.user = user
        self.password = password
        self.dsn = dsn

def _match_group(current_group: str, target_group: str) -> bool:
    return not target_group or current_group == target_group

def _parse_connection(line: str) -> tuple:
    try:
        db_name, conn_str = line.split(':', 1)
        user_pass, dsn = conn_str.split('@', 1)
        user, password = user_pass.split('/', 1)
        return db_name.strip(), user.strip(), password.strip(), dsn.strip()
    except ValueError as e:
        raise ValueError(f"无效连接配置: {line}") from e

def load_alert_config(path: str) -> AlertConfig:
    try:
        with open(path, 'r') as f:
            config = yaml.safe_load(f).get('wx_config', {})
            return AlertConfig(webhook=config.get('webhook'), agentid=config.get('agentid'))
    except FileNotFoundError:
        print(f"警告: 未找到告警配置文件 {path}")
        return None

def parse_dbconfig(path: str, group: str = None) -> Iterator[DBConnectionConfig]:
    current_group = None
    try:
        with open(path, 'r') as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue

                if line.startswith('Group:'):
                    current_group = line.split(' ', 1)[1]
                    if not _match_group(current_group, group):
                        current_group = None
                    continue

                if current_group:
                    db_name, user, password, dsn = _parse_connection(line)
                    yield DBConnectionConfig(db_name, user, password, dsn)

    except FileNotFoundError as e:
        print(f"错误: 未找到数据库配置文件 {e.filename}")
        raise

def send_wxalert(webhook, agentid, content, retry=3):
    for _ in range(retry):
        try:
            response = requests.post(
                webhook,
                json={
                    "msgtype": "markdown",
                    "markdown": {"content": content}
                },
                timeout=5
            )
            response.raise_for_status()
            return True
        except Exception as e:
            print(f"企业微信通知失败: {str(e)}")
    return False

def main():
    import os
    import sqlite3
    start_time = datetime.now()
    

    # 创建数据库目录
    os.makedirs('web_dashboard', exist_ok=True)
    
    # 初始化数据库连接
    db_conn = sqlite3.connect('/home/ai/dev/daycheck/web_dashboard/instance_status.db')
    db_cursor = db_conn.cursor()
    
    
    # 创建状态表
    db_cursor.execute('''
        CREATE TABLE IF NOT EXISTS current_status (
            db_name TEXT PRIMARY KEY,
            status TEXT NOT NULL,
            last_check_time DATETIME NOT NULL,
            details TEXT
        )''')
    db_cursor.execute('''
        CREATE TABLE IF NOT EXISTS status_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            db_name TEXT NOT NULL,
            status TEXT NOT NULL,
            check_time DATETIME NOT NULL,
            details TEXT
        )''')
    db_conn.commit()
    
    # 创建日志目录
    os.makedirs('/home/ai/dev/daycheck/logs', exist_ok=True)
    
    args = parse_args()
    has_error = False
    
    class DualLogger:
        def __init__(self, filename):
            self.terminal = sys.stdout
            self.log = open(filename, 'w')
        
        def write(self, message):
            self.terminal.write(message)
            self.log.write(message)
        
        def flush(self):
            self.terminal.flush()
            self.log.flush()
    
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
    log_path = f'/home/ai/dev/daycheck/logs/check_{timestamp}.log'
    sys.stdout = DualLogger(log_path)
    sys.stderr = sys.stdout
    
    print(f'\n=== 检查开始 {start_time} ===')
        
    for conn in parse_dbconfig(args.config, args.group):
        alert_config = load_alert_config(args.alert_config)
        db_name, user, pwd, dsn = conn.db_name, conn.user, conn.password, conn.dsn
        wx_webhook = alert_config.webhook if alert_config else None
        wx_agentid = alert_config.agentid if alert_config else None
        conn_str = f"{user}/{pwd}@{dsn}"
        result = check_instance(db_name, conn_str)
        print(f"{db_name} : {result}")
            
            # 初始化details变量
        details = str(result) if '失败' in result else ''
        current_status = 'Failed' if '失败' in result else 'Normal'
            
        # 发送企业微信告警
        if current_status == 'Failed':
            # 统一使用f-string格式化
            alert_content = f"""
            <font color="warning">数据库服务异常!</font>
            >服务名称: <font color="comment">{db_name}</font>
            >服务地址：<font color="comment">{dsn}</font>
            >运行状态: <font color="comment">{result}</font>
            >检查时间: <font color="comment">{datetime.now()}</font>
            """
            if wx_webhook and wx_agentid:
                send_wxalert(wx_webhook, wx_agentid, alert_content, args.wx_retry)
            else:
                print("警告：未配置企业微信参数，无法发送告警")

            # 查询最近一次失败记录
            db_cursor.execute('''
                SELECT status FROM current_status 
                WHERE db_name=? 
                ORDER BY last_check_time DESC 
                LIMIT 1
            ''', (db_name,))
            last_record = db_cursor.fetchone()

            # 状态恢复检测逻辑
            if last_record and last_record[0] == 'Failed' and current_status == 'Normal':
                db_cursor.execute('''
                    UPDATE current_status
                    SET status = 'Recovered',
                        details = '自动恢复于: ' || CURRENT_TIMESTAMP
                    WHERE db_name = ? 
                ''', (db_name,))

            # 记录当前状态
            try:
                db_cursor.execute('''
                    INSERT INTO status_history (db_name, status, check_time, details)
                    VALUES (?, ?, ?, ?);
                ''', (db_name, current_status, datetime.now(), details))
                
                db_cursor.execute('''
                    INSERT OR REPLACE INTO current_status 
                    (db_name, status, last_check_time, details)
                    VALUES (?, ?, ?, ?)
                ''', (db_name, current_status, datetime.now(), details))
                db_conn.commit()
            except Exception as e:
                print(f"数据库写入失败: {str(e)}")
                # 发送企业微信通知失败告警
                alert_content = f"状态记录失败!\n数据库: {db_name}\n错误: {str(e)}"
                send_wxalert(wx_webhook, agentid, alert_content, args.wx_retry)
if __name__ == '__main__':
    main()